import uproot
import pandas as pd
sig_name= {"../../Data/data_4tau/signal/THDM_hA_4tau_mh_mch_ATLAS_pt10_14TeV_01.root": 3.67719}
bkg_name= {
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbar.root" : 4452.73,
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarww.root" : 0.0139,
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarz.root" : 0.2136,
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarzz.root" : 0.0001,
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_wwjj.root" : 268.291,
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_zjj.root" : 151060,
        "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_zz.root" : 9.296}


df_sig_list = []
for name,value in sig_name.items():
    df_tmp=uproot.open(name)["datatrain"].arrays(["HT", "MET"],library="pd")
    df_tmp["weight"]=value
    df_sig_list.append(df_tmp)

df_bkg_list =[]
for name,value in bkg_name.items():
    df_tmp=uproot.open(name)["datatrain"].arrays(["HT", "MET"],library="pd")
    df_tmp["weight"]=value
    df_bkg_list.append(df_tmp)
# print(file.keys()) >>> "datatrain"
# print(tree.keys()) >>> "pt", "met" ...
# here are ['weight', 'jet_num', 'signal_jet_num', 'jet_pt1', 'jet_eta1', 'jet_pt2', 'jet_eta2', 'elec_num', 'muon_num', 'lep_num', 'signal_lep_num', 'lep_pt1', 'lep_eta1', 'lep_pt2', 'lep_eta2', 'photon_num', 'MET', 'vis_m', 'vis_pt', 'vis_eta', 'll_inm', 'll_pt', 'll_eta', 'jj_inm', 'jj_pt', 'jj_eta', 'h1_inm', 'h1_pt', 'h1_eta', 'h2_inm', 'h2_pt', 'h2_eta', 'H_inm', 'H_pt', 'H_eta', 'l_l_openangle', 'j_j_openangle', 'l1_j1_openangle', 'l1_j12_openangle', 'l2_j1_openangle', 'l2_j2_openangle', 'h1_h2_openangle', 'HT', 'samesign', 'MT2']


#df = tree.arrays(["h1_inm", "h2_inm"],library="pd")

#df = df[(df!= -10000).all(axis=1)]
df_sig = pd.concat(df_sig_list)
df_bkg = pd.concat(df_bkg_list)



import matplotlib.pyplot as plt

# 假设df_bkg是你的DataFrame
# 选择要绘制直方图的列
column = 'HT'

# 假设你的画布尺寸为 8 英寸宽，6 英寸高
fig, ax = plt.subplots(figsize=(8, 6))

# 绘制直方图
ax.hist(df_bkg["HT"]+df_bkg["MET"], bins=40, range=[0, 600], color='red', density=True, alpha=0.8, histtype='step', label="bkg", linewidth=2)
ax.hist(df_sig["HT"]+df_sig["MET"], bins=40, range=[0, 600], color='black',density=True, alpha=0.8, histtype='step', label="sig", linewidth=2)

# 设置标题和坐标轴标签
ax.set_xlabel("HT [GeV]", ha='right', x=1.0)
ax.tick_params(axis='x', labelsize=20)
x_label = ax.xaxis.get_label()
x_label.set_fontsize(20)
ax.set_ylabel('Normalized Events/15 GeV')
ax.tick_params(axis='y', labelsize=20)
y_label = ax.yaxis.get_label()
y_label.set_fontsize(20)


# 移除图例边框并设置线条样式
handles, labels = ax.get_legend_handles_labels()
leg = ax.legend(handles, labels, fontsize=20, loc='best', frameon=False)

# 确保图例中的线条样式为实线
for line in leg.get_lines():
    line.set_linestyle('-')

# 设置坐标轴参数
ax.tick_params(axis='both', which='both', direction='in', top=True, right=True)

left, bottom, width, height = ax.get_position().bounds
# 调整子图顶部位置，使其上移 1cm
bottom = bottom+0.1
top = bottom+height
left=left+0.05
fig.subplots_adjust(left=left, bottom=bottom, right=left + width, top=top)
# 显示图形
plt.show()
